# -*- coding: utf-8 -*-
import tushare as ts
import pandas as pd
import numpy as np
import datetime
import pymongo
import json
import threading as td
from tushare.util import dateu as du

# 获得股票的分钟级数据 tick -> min
# code: 股票代码
# date: 日期
def get_security_min(code, date):
    str_date = str(date).split(' ')[0];

    # 重试100次
    df = ts.get_tick_data(code,date = str_date, retry_count = 100)

    if df['time'][0] != 'alert("当天没有数据");':

        # 格式化时间
        df['time'] = str_date + ' ' + df['time']
        df['time'] = pd.to_datetime(df['time']);

        # 设置datetime为 index
        df = df.set_index('time');

        # print (df.tail(10));

        price_df = df['price'].resample('1min').ohlc();
        price_df = price_df.dropna();
        # print (price_df.head(10));

        vols = df['volume'].resample('1min').ohlc();
        vols = vols.dropna();
        vol_df = pd.DataFrame(vols, columns = ['volume']);

        # print (vol_df.head(10));

        amounts = df['amount'].resample('1min').ohlc();
        amounts = amounts.dropna();
        amount_df = pd.DataFrame(amounts, columns = ['amount']);

        # print (amount_df.head(10));

        newdf = price_df.merge(vol_df, left_index = True, right_index = True).merge(amount_df, left_index = True, right_index = True);

        newdf['code'] = code;
        newdf['date'] = date;
        newdf['str_date'] = str_date;

        # 最终分钟级数据
        # print (newdf.head(1));

        save_data(newdf, 'stock_minute_data');
        print ('code %s date %s 分钟级数据插入完成' % (code, str_date));

# 获得股票列表
def get_security():

    # 持久化数据
    stock_basics = ts.get_stock_basics();

    # 写入所有股票全部数据
    save_data(stock_basics, 'stock_basics');
    print ('股票列表数据插入完成 get_security success');

    code_array = [];
    date_array = [];

    for index in range(len(stock_basics)):

        if stock_basics['timeToMarket'][index] != 0:
            date_str = str(stock_basics['timeToMarket'][index])[:4] + '-' + str(stock_basics['timeToMarket'][index])[4:6] + '-' + str(stock_basics['timeToMarket'][index])[6:]
            date_str = pd.to_datetime(date_str);
            date_array.append(date_str);
            code_array.append(stock_basics.reset_index()['code'][index]);

    s1 = pd.Series(np.array(code_array));
    s2 = pd.Series(np.array(date_array));
    result_df = pd.DataFrame({
        "code": s1,
        "date": s2
    });
    # print result_df;
    return result_df;

# 获取时间区间
def get_date(start, end):
    date_list = [];
    start = datetime.datetime.strptime(start, '%Y-%m-%d');
    end = datetime.datetime.strptime(end, '%Y-%m-%d');
    end = end + datetime.timedelta(days=1);

    i = datetime.timedelta(days = 0)
    while i < (end - start):
        # date_list.append((start + i).strftime('%Y-%m-%d'));
        date_list.append(pd.to_datetime((start + i).strftime('%Y-%m-%d')));
        i += datetime.timedelta(days = 1);

    print('获取时间列表完成');
    return date_list;

def save_data(df, type):
    client = pymongo.MongoClient('localhost', 27017)
    db = client.quant_data;
    db[type].insert(json.loads(df.reset_index().to_json(orient='records')))

# 保存所有tick 分钟级数据入库
def update_bundle():
    security_list = get_security();
    date_list = get_date('2017-05-19', '2017-05-19');

    # for index in range(len(security_list)):
    for index in range(10):

        # 上市时间
        timeToMarket = security_list['date'][index];
        code = security_list['code'][index]

        for date in date_list:

            # 如果当前时间戳大于股票上市时间
            if date > timeToMarket:

                # 则获取分钟级数据
                get_security_min(code, date);

update_bundle();
